import numpy as np

class MultiTarget:
    def __init__(self, grid_size, num_targets=2):
        self.grid_size = grid_size
        self.num_targets = num_targets
        self.targets = []
        self.speed = 12  # 目标移动速度
        
        # 初始化多个目标
        for i in range(num_targets):
            # 在网格中不同区域生成目标
            if i == 0:
                # 第一个目标在左半部分
                x = np.random.randint(20, grid_size//2 - 20)
                y = np.random.randint(20, grid_size - 20)
            else:
                # 第二个目标在右半部分
                x = np.random.randint(grid_size//2 + 20, grid_size - 20)
                y = np.random.randint(20, grid_size - 20)
            
            target = {
                'id': i,
                'position': (x, y),
                'velocity': (np.random.uniform(-1, 1), np.random.uniform(-1, 1)),
                'is_surrounded': False,
                'escape_attempts': 0,
                'movement_pattern': np.random.choice(['random', 'circular', 'linear']),
                'pattern_timer': 0,
                'pattern_direction': np.random.uniform(0, 2*np.pi)
            }
            self.targets.append(target)
    
    def move_target(self, target_id):
        """移动指定的目标"""
        if target_id >= len(self.targets):
            return
            
        target = self.targets[target_id]
        x, y = target['position']
        
        # 根据移动模式移动目标
        if target['movement_pattern'] == 'random':
            # 随机移动
            dx = np.random.uniform(-self.speed, self.speed)
            dy = np.random.uniform(-self.speed, self.speed)
            
        elif target['movement_pattern'] == 'circular':
            # 圆形移动
            target['pattern_timer'] += 0.1
            radius = 30
            center_x = self.grid_size // 2 if target_id == 0 else self.grid_size * 3 // 4
            center_y = self.grid_size // 2
            
            new_x = center_x + radius * np.cos(target['pattern_timer'])
            new_y = center_y + radius * np.sin(target['pattern_timer'])
            dx = new_x - x
            dy = new_y - y
            
        else:  # linear
            # 直线移动，碰到边界反弹
            dx = self.speed * np.cos(target['pattern_direction'])
            dy = self.speed * np.sin(target['pattern_direction'])
            
            # 检查边界反弹
            if x + dx <= 10 or x + dx >= self.grid_size - 10:
                target['pattern_direction'] = np.pi - target['pattern_direction']
                dx = -dx
            if y + dy <= 10 or y + dy >= self.grid_size - 10:
                target['pattern_direction'] = -target['pattern_direction']
                dy = -dy
        
        # 更新位置
        new_x = max(10, min(x + dx, self.grid_size - 10))
        new_y = max(10, min(y + dy, self.grid_size - 10))
        
        target['position'] = (new_x, new_y)
        target['velocity'] = (dx, dy)
    
    def try_escape(self, target_id, uav_positions):
        """目标尝试逃脱合围"""
        if target_id >= len(self.targets):
            return
            
        target = self.targets[target_id]
        x, y = target['position']
        
        # 计算远离所有无人机的方向
        escape_direction = np.array([0.0, 0.0])
        for uav_pos in uav_positions:
            direction_away = np.array([x - uav_pos[0], y - uav_pos[1]])
            distance = np.linalg.norm(direction_away)
            if distance > 0:
                # 距离越近，逃脱力度越大
                weight = 1.0 / (distance + 1)
                escape_direction += weight * direction_away / distance
                
        # 归一化逃脱方向
        if np.linalg.norm(escape_direction) > 0:
            escape_direction = escape_direction / np.linalg.norm(escape_direction)
            
        # 按逃脱方向移动，速度更快
        escape_speed = self.speed * 1.5
        new_x = x + escape_direction[0] * escape_speed
        new_y = y + escape_direction[1] * escape_speed
        
        # 边界检查
        new_x = max(10, min(new_x, self.grid_size - 10))
        new_y = max(10, min(new_y, self.grid_size - 10))
        
        target['position'] = (new_x, new_y)
        target['escape_attempts'] += 1
    
    def move_all_targets(self):
        """移动所有目标"""
        for i in range(len(self.targets)):
            self.move_target(i)
    
    def get_target_position(self, target_id):
        """获取指定目标的位置"""
        if target_id < len(self.targets):
            return self.targets[target_id]['position']
        return None
    
    def get_all_positions(self):
        """获取所有目标的位置"""
        return [target['position'] for target in self.targets]
    
    def get_state(self):
        """获取所有目标的状态"""
        state = []
        for target in self.targets:
            state.extend([target['position'][0], target['position'][1]])
        return np.array(state)
    
    def reset(self):
        """重置所有目标"""
        for i, target in enumerate(self.targets):
            # 重新随机化位置
            if i == 0:
                x = np.random.randint(20, self.grid_size//2 - 20)
                y = np.random.randint(20, self.grid_size - 20)
            else:
                x = np.random.randint(self.grid_size//2 + 20, self.grid_size - 20)
                y = np.random.randint(20, self.grid_size - 20)
            
            target['position'] = (x, y)
            target['velocity'] = (np.random.uniform(-1, 1), np.random.uniform(-1, 1))
            target['is_surrounded'] = False
            target['escape_attempts'] = 0
            target['movement_pattern'] = np.random.choice(['random', 'circular', 'linear'])
            target['pattern_timer'] = 0
            target['pattern_direction'] = np.random.uniform(0, 2*np.pi)